//
// Created by 34753 on 2025/5/19.
//

#ifndef FFT_H
#define FFT_H
#include <cstring>
#include<math.h>

#ifdef __ARM__
#include <bits/move.h>
#endif

namespace utils
{
    class FFT
    {
    public:
        template<unsigned int N>
        static void f32(float dataR[], float dataI[])
        {
            // 1. 位反转置换（Bit-reversal permutation） 由于分支预测过多，会导致分支惩罚比较严重
            const unsigned int M = log2(N); // 需确保N是2的幂
            for (unsigned int i = 0; i < N; ++i)
            {
                // 计算i的位反转j
                unsigned int j = 0;
                for (unsigned int k = 0; k < M; ++k)
                {
                    j = (j << 1) | ((i >> k) & 1);
                }
                if (j > i)
                {
                    std::swap(dataR[i], dataR[j]);
                    std::swap(dataI[i], dataI[j]);
                }
            }

            // 2. 蝶形运算（Butterfly Computation）
            for (int m = 1; m < N; m <<= 1)// 层级循环
            {
                // m为当前阶段蝶形间距（1,2,4...64）
                const int m2 = m * 2; // 每组数据长度
                for (int k = 0; k < m; k++)//  组内循环，奇偶组的长度
                {
                    // 遍历每组内的旋转因子索引k
                    // 计算旋转因子W_{m2}^k = e^{-j2πk/m2}
                    const float theta = -2.0 * M_PI * k / m2;
                    const float wR = cos(theta); // 实部
                    const float wI = sin(theta); // 虚部

                    // 遍历所有相同旋转因子的蝶形组，步长m2（其实是单个奇偶组的跨度）
                    for (int i = k, j = 0; i < N; i += m2)// 奇偶组的数量
                    {
                        // 复数乘法
                        j = i + m; // 蝶形对第二个元素的索引 此时的m就是N/2 在组内变换
                        // 计算旋转因子与data[j]的乘积
                        const float tr = wR * dataR[j] - wI * dataI[j];
                        const float ti = wR * dataI[j] + wI * dataR[j];
                        // 蝶形运算：更新当前对的两个元素
                        dataR[j] = dataR[i] - tr; // 下半部分
                        dataI[j] = dataI[i] - ti;
                        dataR[i] += tr; // 上半部分
                        dataI[i] += ti;
                    }
                }
            }
        }

        /**
         * @brief 执行128点快速傅里叶变换（FFT）的蝶形算法实现。
         *
         * @details
         * 该函数对输入的实部和虚部数组进行FFT计算，输出频域复数结果覆盖输入数组。
         * 实现了原位（in-place）计算的Cooley-Tukey算法，包含位反转置换和7级蝶形运算。
         *
         * @param[in,out] dataR 输入实部数组（长度至少128），FFT结果将覆盖此数组。
         * @param[in,out] dataI 输入虚部数组（长度至少128），FFT结果将覆盖此数组。
         * ### 算法流程
         * 1. **位反转置换（Bit-reversal Permutation）**
         *    - 将输入索引的二进制位反转后重新排序，确保FFT正确性。
         *    - 例如，原索引6（二进制`0000110`）变为`0110000`（十进制48）。
         *    - 使用dataI作为临时存储，避免数据覆盖。
         *
         * 2. **预计算旋转因子**
         *    - 复数旋转因子 $ W_N^k = e^{-2\pi jk/N} $ 的实部（cos）和虚部（-sin）已编码于表中。
         *    - 索引表长度64（即 $ N/2 = 64 $），避免运行时重复计算三角函数。
         *
         * 3. **蝶形运算（七级迭代计算）**
         *    - **级数**：$ L = 1,2,\dots,7 $，对应跨度为 $ b = 2^{L-1} $。
         *    - **蝶形组处理**：
         *      - 计算旋转因子索引 $ p = j \cdot 2^{7-L} $（$ j $为组索引）
         *      - 每对元素 $ (k, k+b) $ 进行复数乘加操作：
         *        $$
         *        X[k] \leftarrow X[k] + W \cdot X[k+b] \\
         *        X[k+b] \leftarrow X[k] - W \cdot X[k+b]
         *
        $$
         *        其中 $ W $ 是旋转因子表中的复数。
         *
         * @note
         * - **输入要求**：dataR和dataI必须是长度≥128的非空数组。
         * - **原位计算**：直接覆盖输入数组，调用前需保存原数据（如果需保留）。
         * - **输出格式**：FFT结果以复数形式分布在dataR（实部）和dataI（虚部）中。
         * - **扩展性**：硬编码128点FFT，不支持其他点数配置。
         */
        static void f32_128(float dataR[], float dataI[])
        {
            unsigned short x1, x2, x3, x4, x5, x6, i;
            // 1. 位反转置换（Bit-reversal permutation） 将输入数据的索引按二进制位反转顺序重新排列，实现原位FFT的输入准备
            for (i = 0; i < 128; i++)
            {
                // 将i分解为7位二进制（因为128=2^7），存储每一位的值
                unsigned short x0 = x1 = x2 = x3 = x4 = x5 = x6 = 0;
                x0 = i & 0x01; // 最低位（2^0位）
                x1 = (i / 2) & 0x01; // 2^1位
                x2 = (i / 4) & 0x01; // 2^2位
                x3 = (i / 8) & 0x01; // 2^3位
                x4 = (i / 16) & 0x01; // 2^4位
                x5 = (i / 32) & 0x01; // 2^5位
                x6 = (i / 64) & 0x01; // 2^6位（最高位）

                // 计算位反转后的索引：原最高位变为最低位，依此类推
                const unsigned short xx = x0 * 64 + x1 * 32 + x2 * 16 + x3 * 8 + x4 * 4 + x5 * 2 + x6;
                dataI[xx] = dataR[i]; // 将原数据按反转索引存入虚部数组
            }
            std::memcpy(dataR, dataI, sizeof(float) * 128);
            std::memset(dataI, 0, sizeof(float) * 128);

            // 预计算旋转因子表
            const float cos_tab[64] = {
                1.000000f, 0.998795f, 0.995185f, 0.989177f, 0.980785f, 0.970031f, 0.956940f, 0.941544f, 0.923880f,0.903989f, 0.881921f, 0.857729f, 0.831470f, 0.803208f, 0.773010f, 0.740951f,
                0.707107f, 0.671559f, 0.634393f, 0.595699f, 0.555570f, 0.514103f, 0.471397f, 0.427555f, 0.382683f,0.336890f, 0.290285f, 0.242980f, 0.195090f, 0.146730f, 0.098017f, 0.049068f,
                -0.000000f, -0.049068f, -0.098017f, -0.146730f, -0.195090f, -0.242980f, -0.290285f, -0.336890f,-0.382683f, -0.427555f, -0.471397f, -0.514103f, -0.555570f, -0.595699f, -0.634393f, -0.671559f,
                -0.707107f, -0.740951f, -0.773010f, -0.803208f, -0.831470f, -0.857729f, -0.881921f, -0.903989f,-0.923880f, -0.941544f, -0.956940f, -0.970031f, -0.980785f, -0.989177f, -0.995185f, -0.998795f,
            };

            const float sin_tab[64] = {
                -0.000000f, -0.049068f, -0.098017f, -0.146730f, -0.195090f, -0.242980f, -0.290285f, -0.336890f,-0.382683f, -0.427555f, -0.471397f, -0.514103f, -0.555570f, -0.595699f, -0.634393f, -0.671559f,
                -0.707107f, -0.740951f, -0.773010f, -0.803208f, -0.831470f, -0.857729f, -0.881921f, -0.903989f,-0.923880f, -0.941544f, -0.956940f, -0.970031f, -0.980785f, -0.989177f, -0.995185f, -0.998795f,
                -1.000000f, -0.998795f, -0.995185f, -0.989177f, -0.980785f, -0.970031f, -0.956940f, -0.941544f,-0.923880f, -0.903989f, -0.881921f, -0.857729f, -0.831470f, -0.803208f, -0.773010f, -0.740951f,
                -0.707107f, -0.671559f, -0.634393f, -0.595699f, -0.555570f, -0.514103f, -0.471397f, -0.427555f,-0.382683f, -0.336890f, -0.290285f, -0.242980f, -0.195090f, -0.146731f, -0.098017f, -0.049068f,
            };

            // 2. 蝶形运算（Butterfly Computation） 进行7级蝶形运算（128点FFT需要log2(128)=7级）
            for (unsigned short L = 1; L <= 7; L++)
            {
                /* 计算当前级数对应的蝶形跨度b=2^(L-1) */
                unsigned short b = 1;
                i = L - 1;
                while (i > 0)
                {
                    b *= 2; // 每次循环乘以2，最终得到b=2^(L-1)
                    --i;
                }

                // 遍历每个蝶形组（共b个组）
                for (unsigned short j = 0; j <= b - 1; j++)
                {
                    /* 计算旋转因子索引p = (2^(7-L)) * j */
                    unsigned short p = 1;
                    i = 7 - L; // 计算指数部分
                    while (i > 0) // 计算2^(7-L)
                    {
                        p *= 2;
                        --i;
                    }
                    p *= j; // 最终p = j * 2^(7-L)

                    // 遍历当前组内的蝶形运算元素对 步长为2b，处理每个蝶形对
                    for (int k = j; k < 128; k += 2 * b)
                    {
                        const float TR = dataR[k];// 保存当前元素k的实部和虚部
                        const float TI = dataI[k];

                        const float temp = dataR[k + b];// 临时保存k+b元素的实部（用于后续计算）

                        // 计算旋转因子乘法的结果（复数乘法） W = cos_tab[p] + i*sin_tab[p]
                        dataR[k] = TR + dataR[k + b] * cos_tab[p] + dataI[k + b] * sin_tab[p];// X[k+b] * W 的实部：dataR[k+b]*cos - dataI[k+b]*sin
                        dataI[k] = TI - dataR[k + b] * sin_tab[p] + dataI[k + b] * cos_tab[p];// X[k+b] * W 的虚部：dataR[k+b]*sin + dataI[k+b]*cos

                        // 更新元素k+b的值：X[k+b] = X[k] - X[k+b] * W
                        dataR[k + b] = TR - dataR[k + b] * cos_tab[p] - dataI[k + b] * sin_tab[p];
                        dataI[k + b] = TI + temp * sin_tab[p] - dataI[k + b] * cos_tab[p];
                    }
                }
            }
        }

        template<unsigned short N>
        static void f32_spectrum(float dataR[], float dataI[], float spectrum[])
        {
            for (unsigned short i = 0; i < N; ++i)
            {
                // 修改为归一化
                spectrum[i] = sqrt(dataR[i] * dataR[i] + dataI[i] * dataI[i]) / (N / 2); // 对于非直流分量除以N/2
            }
            spectrum[0] /= 2; //  直流分量除以N
        }
    };
}


#endif //FFT_H
